"""
预测结果分析脚本

该脚本用于分析指定日期范围内的预测文件，并按composite_score排名分段计算B1_tag为1的概率。

功能:
1. 自动扫描predictions目录中的所有预测文件
2. 支持指定日期范围分析
3. 按composite_score排名分段(前10%, 10-20%, ..., 90-100%)
4. 计算每个分段中B1_tag为1的概率
5. 输出结果为CSV格式文件

使用方法:
1. 分析所有日期范围的数据:
   python analyze_predictions_by_date_range.py

2. 指定日期范围:
   python analyze_predictions_by_date_range.py --start-date 2025-03-01 --end-date 2025-04-01

3. 指定输入目录:
   python analyze_predictions_by_date_range.py --input-dir /path/to/predictions

4. 指定输出文件名:
   python analyze_predictions_by_date_range.py --output-file my_analysis_result.csv

5. 组合使用:
   python analyze_predictions_by_date_range.py --start-date 2025-03-01 --end-date 2025-04-01 --input-dir /path/to/predictions --output-file result.csv
"""

import pandas as pd
import os
import argparse
from datetime import datetime
import numpy as np

def get_date_range_from_files(predictions_dir):
    """
    从预测文件名中提取所有日期，并返回日期范围
    """
    files = [f for f in os.listdir(predictions_dir) if f.startswith('predictions_') and f.endswith('.csv')]
    dates = []
    for file in files:
        try:
            # 从文件名中提取日期，例如 predictions_20250317.csv -> 20250317
            date_str = file.split('_')[1].split('.')[0]
            date = datetime.strptime(date_str, '%Y%m%d')
            dates.append(date)
        except Exception as e:
            print(f"无法解析文件 {file} 的日期: {e}")
    
    if not dates:
        return None, None
        
    return min(dates), max(dates)

def load_predictions_in_date_range(predictions_dir, start_date, end_date):
    """
    加载指定日期范围内的所有预测数据
    
    Parameters:
    predictions_dir: str, 预测文件目录路径
    start_date: datetime, 开始日期
    end_date: datetime, 结束日期
    
    Returns:
    DataFrame: 所有预测数据的合并
    """
    all_predictions = []
    
    # 获取所有预测文件
    files = [f for f in os.listdir(predictions_dir) if f.startswith('predictions_') and f.endswith('.csv')]
    
    for file in files:
        try:
            # 从文件名中提取日期
            date_str = file.split('_')[1].split('.')[0]
            file_date = datetime.strptime(date_str, '%Y%m%d')
            
            # 检查日期是否在指定范围内
            if start_date <= file_date <= end_date:
                file_path = os.path.join(predictions_dir, file)
                df = pd.read_csv(file_path)
                # 添加日期列（如果不存在）
                if 'date' not in df.columns:
                    df['date'] = file_date.strftime('%Y-%m-%d')
                all_predictions.append(df)
        except Exception as e:
            print(f"处理文件 {file} 时出错: {e}")
    
    if not all_predictions:
        return pd.DataFrame()
    
    # 合并所有数据
    combined_df = pd.concat(all_predictions, ignore_index=True)
    return combined_df

def analyze_by_score_percentiles(df, score_column='composite_score'):
    """
    按照分数排名分段分析B1_tag为1的概率
    
    Parameters:
    df: DataFrame, 包含预测数据的DataFrame
    score_column: str, 用于分档的分数列名
    
    Returns:
    DataFrame: 各分段的统计结果
    """
    # 按分数从高到低排序
    df_sorted = df.sort_values(by=score_column, ascending=False).reset_index(drop=True)
    
    # 计算每个分段的边界
    total_count = len(df_sorted)
    segment_sizes = [int(total_count * 0.1 * (i + 1)) for i in range(10)]
    
    # 初始化结果列表
    results = []
    
    # 定义分段标签
    segment_labels = [f"前{i*10}-{(i+1)*10}%" for i in range(10)]
    
    start_idx = 0
    for i, end_idx in enumerate(segment_sizes):
        # 获取当前分段的数据
        if i == 0:
            segment_data = df_sorted.iloc[start_idx:end_idx]
        else:
            segment_data = df_sorted.iloc[start_idx:end_idx]
            
        # 计算统计数据
        count = len(segment_data)
        b1_tag_count = (segment_data['B1_tag'] == 1).sum()
        b1_tag_probability = b1_tag_count / count if count > 0 else 0
        
        # 获取分数范围
        if len(segment_data) > 0:
            min_score = segment_data[score_column].min()
            max_score = segment_data[score_column].max()
            score_range = f"{min_score:.4f} - {max_score:.4f}"
        else:
            score_range = "N/A"
        
        # 添加到结果
        results.append({
            'score_segment': segment_labels[i],
            '总数量': count,
            'B1_tag为1的数量': b1_tag_count,
            'B1_tag为1的概率': round(b1_tag_probability, 4),
            '分数范围': score_range
        })
        
        # 更新起始索引
        start_idx = end_idx
    
    # 转换为DataFrame
    result_df = pd.DataFrame(results)
    return result_df

def main():
    parser = argparse.ArgumentParser(description='分析指定日期范围内的预测文件')
    parser.add_argument('--start-date', '--start_date', type=str, help='开始日期 (YYYY-MM-DD)')
    parser.add_argument('--end-date', '--end_date', type=str, help='结束日期 (YYYY-MM-DD)')
    parser.add_argument('--input-dir', '--input_dir', type=str, help='输入目录路径 (默认: predictions)')
    parser.add_argument('--output-file', '--output_file', type=str, default='prediction_analysis_result.csv', 
                        help='输出文件名 (默认: prediction_analysis_result.csv)')
    
    args = parser.parse_args()
    
    # 预测文件目录
    if args.input_dir:
        predictions_dir = args.input_dir
    else:
        predictions_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'predictions')
    
    # 检查目录是否存在
    if not os.path.exists(predictions_dir):
        print(f"输入目录不存在: {predictions_dir}")
        return
    
    # 获取日期范围
    min_date, max_date = get_date_range_from_files(predictions_dir)
    
    if min_date is None or max_date is None:
        print("未找到有效的预测文件")
        return
    
    print(f"可用日期范围: {min_date.strftime('%Y-%m-%d')} 到 {max_date.strftime('%Y-%m-%d')}")
    
    # 确定实际使用的日期范围
    if args.start_date:
        try:
            start_date = datetime.strptime(args.start_date, '%Y-%m-%d')
        except ValueError:
            print("开始日期格式错误，请使用 YYYY-MM-DD 格式")
            return
    else:
        start_date = min_date
    
    if args.end_date:
        try:
            end_date = datetime.strptime(args.end_date, '%Y-%m-%d')
        except ValueError:
            print("结束日期格式错误，请使用 YYYY-MM-DD 格式")
            return
    else:
        end_date = max_date
    
    # 确保日期范围有效
    if start_date > end_date:
        print("开始日期不能晚于结束日期")
        return
    
    print(f"分析日期范围: {start_date.strftime('%Y-%m-%d')} 到 {end_date.strftime('%Y-%m-%d')}")
    
    # 加载指定日期范围内的数据
    print("正在加载预测数据...")
    df = load_predictions_in_date_range(predictions_dir, start_date, end_date)
    
    if df.empty:
        print("在指定日期范围内未找到数据")
        return
    
    print(f"共加载 {len(df)} 条预测记录")
    
    # 按composite_score排名分段分析
    print("正在按composite_score排名分段分析...")
    composite_result = analyze_by_score_percentiles(df, 'composite_score')
    
    # 添加列标识分析类型
    composite_result.insert(0, '分析类型', '综合评分')
    
    # 对其他评分进行分析
    print("正在按rf_prediction排名分段分析...")
    rf_result = analyze_by_score_percentiles(df, 'rf_prediction')
    rf_result.insert(0, '分析类型', '随机森林')
    
    print("正在按lr_prediction排名分段分析...")
    lr_result = analyze_by_score_percentiles(df, 'lr_prediction')
    lr_result.insert(0, '分析类型', '逻辑回归')
    
    print("正在按svm_prediction排名分段分析...")
    svm_result = analyze_by_score_percentiles(df, 'svm_prediction')
    svm_result.insert(0, '分析类型', '支持向量机')
    
    print("正在按nn_prediction排名分段分析...")
    nn_result = analyze_by_score_percentiles(df, 'nn_prediction')
    nn_result.insert(0, '分析类型', '神经网络')
    
    # 合并所有结果
    all_results = pd.concat([composite_result, rf_result, lr_result, svm_result, nn_result], ignore_index=True)
    
    # 保存结果
    output_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), args.output_file)
    all_results.to_csv(output_path, index=False, encoding='utf-8-sig')
    
    print(f"分析结果已保存到: {output_path}")
    print("\n分析结果:")
    print(all_results.to_string(index=False))

if __name__ == "__main__":
    main()